from core.function import Function
from core.cuda import cuda_module


class Softmax(Function):
    def __init__(self, axis=1):
        """
        初始化 Softmax 操作的函数。

        Args:
            axis (int, optional): Softmax 操作的轴。默认为1。
        """
        self.axis = axis

    def forward(self, x):
        """
        计算 Softmax 操作的前向传播。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: Softmax 操作的输出变量。
        """
        xp = cuda_module
        y = x - x.max(axis=self.axis, keepdims=True)
        y = xp.exp(y)
        y /= y.sum(axis=self.axis, keepdims=True)
        return y

    def backward(self, gy):
        """
        计算 Softmax 操作的反向传播。

        Args:
            gy (Variable): 输出的梯度。

        Returns:
            Variable: 输入的梯度。
        """
        y = self.outputs[0]()
        gx = y * gy
        sumdx = gx.sum(axis=self.axis, keepdims=True)
        gx = gx - y * sumdx
        return gx

    def softmax(self, x):
        """
        对输入变量应用 Softmax 操作。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: Softmax 操作的输出变量。
        """
        return self(x)
